(* Written by Erica Sponsler and Nate Weiss *)

open Ast
open Bytecode

(* Stack layout just after "Ent":

              <-- SP
   Local n
   ...
   Local 0
   Saved FP   <-- FP
   Saved PC
   Arg 0
   ...
   Arg n *)
module NodeMap = Map.Make(String)

type node = {
  value: int;
  children: int list;
  }

type env = {
    stack: int array;
    node_heap: (node) NodeMap.t;
    globals: int array;
    saved_sp: int array;
  }

let rec modify_ith_element new_val i target list = 
  if i==target then (match list with 
    [] -> new_val :: []
  | hd::tl -> new_val :: tl)
      else (match list with
	hd::tl -> hd :: (modify_ith_element new_val (i+1) target tl)
      | _ -> raise (Failure ("Invalid List")))

let execute_prog prog =

  let rec exec fp sp hp sc pc env = match prog.text.(pc) with
(*    Lit i  -> stack.(sp) <- i ; exec fp (sp+1) (pc+1)*)
    LitI i -> env.stack.(sp) <- i ; exec fp (sp+1) hp sc (pc+1) env
  | LitC c -> env.stack.(sp) <- (int_of_char c) ; exec fp (sp+1) hp sc (pc+1) env
  | LitB b -> env.stack.(sp) <- if b then 1 else 0 ; exec fp (sp+1) hp sc (pc+1) env
  | LitNull -> env.stack.(sp) <- (-1) ; exec fp (sp+1) hp sc (pc+1) env
(*  | Drp -> exec fp (sp-1) (pc+1) *)
  | Drp i -> exec fp (sp-1) hp sc (pc+1) env
  | Bin op -> let op1 = env.stack.(sp-2) and op2 = env.stack.(sp-1) in
      env.stack.(sp-2) <- (let boolean i = if i then 1 else 0 in
      let b1 = (op1 == 1) in
      let b2 = (op2 == 1) in
      match op with
	Add     -> op1 + op2
      | Sub     -> op1 - op2
      | Mult    -> op1 * op2
      | Div     -> op1 / op2
      | Equal   -> boolean (op1 =  op2)
      | Neq     -> boolean (op1 != op2)
      | Less    -> boolean (op1 <  op2)
      | Leq     -> boolean (op1 <= op2)
      | Greater -> boolean (op1 >  op2)
      | Geq     -> boolean (op1 >= op2)
      | BoolAnd -> boolean (b1 && b2)
      | BoolOr  -> boolean (b1 || b2)) ;
      exec fp (sp-1) hp sc (pc+1) env
  | Lod -> env.stack.(sp-2) <- env.globals.(env.stack.(sp-2)) ; exec fp (sp-1) hp sc (pc+1) env
  | Str ->  env.globals.(env.stack.(sp-1)) <- env.stack.(sp-2) ; exec fp (sp-1) hp sc (pc+1) env
  | Lfp -> env.stack.(sp-2) <- env.stack.(fp+env.stack.(sp-2)) ; exec fp (sp-1) hp sc (pc+1) env
  | Sfp -> env.stack.(fp+env.stack.(sp-1)) <- env.stack.(sp-2) ; exec fp (sp-1) hp sc (pc+1) env
  | Ldh -> (try let node_val = (NodeMap.find (string_of_int env.stack.(sp-2)) env.node_heap) in (match env.stack.(sp-1) with
      (-1) -> env.stack.(sp-2) <- node_val.value ; exec fp (sp-1) hp sc (pc+1) env
    | _ -> env.stack.(sp-2) <- (List.nth node_val.children env.stack.(sp-1)) ; exec fp (sp-1) hp sc (pc+1) env) with Not_found -> raise (Failure ("Failed to find node in Ldh")))
  | Sth ->  (try let node_val = (NodeMap.find (string_of_int env.stack.(sp-2)) env.node_heap) in (match env.stack.(sp-1) with
      (-1) -> let new_heap = (NodeMap.add (string_of_int env.stack.(sp-2)) {node_val with value = env.stack.(sp-3)} env.node_heap) in exec fp (sp-2) hp sc (pc+1) {env with node_heap = new_heap}
    | _ -> let new_node = {node_val with children = (modify_ith_element env.stack.(sp-3) 0 env.stack.(sp-1) node_val.children)} in let new_heap = (NodeMap.add (string_of_int env.stack.(sp-2)) new_node env.node_heap) in exec fp (sp-2) hp sc (pc+1) {env with node_heap = new_heap}) with Not_found -> raise (Failure ("Failed to find node in Sth")))
  | Cnd -> let new_heap = (NodeMap.add (string_of_int hp) {value = env.stack.(sp-1); children = []} env.node_heap) in
    env.stack.(sp-1) <- hp ; exec fp sp (hp+1) sc (pc+1) {env with node_heap = new_heap}
  | Jsr(-1) -> print_endline (string_of_int env.stack.(sp-1)) ; exec fp sp hp sc (pc+1) env
  | Jsr(-2) -> print_endline (if (env.stack.(sp-1) == 1) then "true" else "false") ; exec fp sp hp sc (pc+1) env
  | Jsr(-3) -> print_endline (Char.escaped (char_of_int env.stack.(sp-1))) ; exec fp sp hp sc (pc+1) env
  | Jsr i   -> env.stack.(sp)   <- pc + 1       ; exec fp (sp+1) hp sc i env
  | Ssp -> env.saved_sp.(sc) <- sp               ; exec fp (sp + env.stack.(sp-1) - 1) hp (sc+1) (pc+1) env
  | Rsp ->  let new_sp = env.saved_sp.(sc-1) in exec fp new_sp hp (sc-1) (pc+1) env
  | Ent i   -> env.stack.(sp)   <- fp           ; exec sp (sp+i+1) hp sc (pc+1) env
  | Rts i   -> let new_fp = env.stack.(fp) and new_pc = env.stack.(fp-1) in
               env.stack.(fp-i-1) <- env.stack.(sp-1) ; exec new_fp (fp-i) hp sc new_pc env
  | Beq i   -> exec fp (sp-1) hp sc (pc + if env.stack.(sp-1) =  0 then i else 1) env
  | Bne i   -> exec fp (sp-1) hp sc (pc + if env.stack.(sp-1) != 0 then i else 1) env
  | Bra i   -> exec fp sp hp sc (pc+i) env
  | Hlt     -> ()

  in let env = { stack = Array.make 1024 (-1);
		 node_heap = NodeMap.empty;
		 globals = Array.make prog.num_globals (-1);
		 saved_sp = Array.make 1024 (-1) }

  in exec 0 0 0 0 0 env
